import os
from sklearn.model_selection import train_test_split
from  sklearn.model_selection import  train_test_split
# 定义源文件夹路径
image_folder = r'D:\xunlian\custom_training\JPEGImages'
label_folder = r'D:\xunlian\custom_training\labels'

# 获取所有图片文件名（假设都是jpg格式）
image_files = os.listdir(image_folder)

# 拆分数据集为训练集和验证集 (例如 80% 训练集, 20% 验证集)
train_files, val_files = train_test_split(image_files, test_size=0.2, random_state=42)
print(train_files)
# 定义输出 txt 文件路径
train_txt_path = r'D:\xunlian\custom_training\train.txt'
val_txt_path = r'D:\xunlian\custom_training\val.txt'

# 将训练集图片路径写入 train.txt
# 打开训练集文本文件进行写入
train_txt = open(train_txt_path, 'w')
for image_file in train_files:
    image_path = os.path.join(image_folder, image_file)
    train_txt.write(f"{image_path}\n")
train_txt.close()  # 记得关闭文件

# 打开验证集文本文件进行写入
val_txt = open(val_txt_path, 'w')
for image_file in val_files:
    image_path = os.path.join(image_folder, image_file)
    val_txt.write(f"{image_path}\n")
val_txt.close()  # 记得关闭文件

print("数据集拆分完成，并生成了 train.txt 和 val.txt 文件！")